import os
import pandas as pd
import time
import argparse
import seaborn as sn
import numpy as np
import matplotlib.pyplot as plt
# sn.set(style="darkgrid") #这是seaborn默认的风格


from model import SASRec
from utils import *



def str2bool(s):
    if s not in {'false', 'true'}:
        raise ValueError('Not a valid boolean string')
    return s == 'true'
def auc_print(x_data,y_data):
    # x_data = pd.DataFrame(x_data)
    # y_data = pd.DataFrame(y_data)
    # objs = [x_data, y_data]
    # data = pd.concat(objs, axis=0, join='outer', ignore_index=False, copy=True)
    ax = sn.lineplot(x=x_data, y=y_data, )
    plt.show()
parser = argparse.ArgumentParser()
parser.add_argument('--dataset',default='trainL.csv') # 完整数据集导入
parser.add_argument('--user_train',default="user_trainL.csv")
parser.add_argument('--user_test',default="user_testL.csv")
# parser.add_argument('--train_dir', default="cuda")
parser.add_argument('--train_dir', default="cpu")
# 使用cpu训练
parser.add_argument('--batch_size', default=1000, type=int)
# 批尺寸
parser.add_argument('--lr', default=0.01, type=float)
# Adam分类器lr
parser.add_argument('--hidden_units', default=10, type=int)
# 词向量维度(自注意力层）
parser.add_argument('--hidden_units_FM', default=10, type=int)
# FM隐向量维度
parser.add_argument('--latent_units', default=5, type=int)
parser.add_argument('--num_blocks', default=1, type=int)
parser.add_argument('--num_epochs', default=20, type=int)
# 迭代次数
parser.add_argument('--num_heads', default=4, type=int)
# 多头注意力
parser.add_argument('--dropout_rate', default=0.2, type=float)
# 丢弃率（dropout）
parser.add_argument('--l2_emb', default=0.0005, type=float)
#parser.add_argument('--device', default='cuda', type=str) # GPU训练
parser.add_argument('--device', default='cpu', type=str) # cpu训练
parser.add_argument('--state_dict_path', default=None, type=str)
parser.add_argument('--FMmode', default="FM", type=str)   # ['FM','AFM','MFM','HoAFM', 'TransFM','ATransFM',"None"]
# 因子分解机层切换
parser.add_argument('--Biasmode', default='None', type=str)  # ['lognormal', 'normal', 'abs', 'learn', 'None']

t_auc = []
t_acc = []
t_epoch = []

args = parser.parse_args()
if not os.path.isdir(args.dataset + '_' + args.train_dir):
    os.makedirs(args.dataset + '_' + args.train_dir)
with open(os.path.join(args.dataset + '_' + args.train_dir, 'args.txt'), 'w') as f:
    f.write('\n'.join([str(k) + ',' + str(v) for k, v in sorted(vars(args).items(), key=lambda x: x[0])]))
f.close()

if __name__ == '__main__':
    # global dataset
    dataset = data_partition(args.user_train,args.user_test) # utils-date_partition

    [user_train, user_test] = dataset
    # emb参数
    emb_list = [len(set(user_train[:,i])) for i in range(1,user_train.shape[1])]
    print(emb_list) #

    parser.add_argument('--emb_list', default=emb_list, type=list)
    parser.add_argument('--maxlen', default=len(emb_list), type=int)
    parser.add_argument('--maxlen_FM', default=sum(emb_list), type=int)
    args = parser.parse_args()

    num_batch = len(user_train) // args.batch_size # tail? + ((len(user_train) % args.batch_size) != 0)
    
    f = open(os.path.join(args.dataset + '_' + args.train_dir, 'log.txt'), 'w')
    sampler = WarpSampler(user_train, batch_size=args.batch_size, n_workers=1)

    model = SASRec(args).to(args.device) # no ReLU activation in original SASRec implementation?
    # model-SASRec
    for name, param in model.named_parameters():
        try:
            torch.nn.init.xavier_normal_(param.data)
        except:
            pass # just ignore those failed init layers
    
    # this fails embedding init 'Embedding' object has no attribute 'dim'
    # model.apply(torch.nn.init.xavier_uniform_)

    model.train() # enable model training


    epoch_start_idx = 1
    if args.state_dict_path is not None:
        try:
            model.load_state_dict(torch.load(args.state_dict_path, map_location=torch.device(args.device)))
            tail = args.state_dict_path[args.state_dict_path.find('epoch=') + 6:]
            epoch_start_idx = int(tail[:tail.find('.')]) + 1
        except: # in case your pytorch version is not 1.6 etc., pls debug by pdb if load weights failed
            print('failed loading state_dicts, pls check file path: ', end="")
            print(args.state_dict_path)
            print('pdb enabled for your quick check, pls type exit() if you do not need it')
            import pdb; pdb.set_trace()


        # ce_criterion = torch.nn.CrossEntropyLoss()
    # https://github.com/NVIDIA/pix2pixHD/issues/9 how could an old bug appear again...
    bce_criterion = torch.nn.BCEWithLogitsLoss().to(args.device) # pos_weight=torch.FloatTensor([1])
    adam_optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, betas=(0.9, 0.98))
    
    T = 0.0
    t0 = time.time()
    
    for epoch in range(epoch_start_idx, args.num_epochs + 1):
        for step in range(num_batch): # tqdm(range(num_batch), total=num_batch, ncols=70, leave=False, unit='b'):
            seq, pos = sampler.next_batch() # tuples to ndarray
            seq, pos = np.array(seq), torch.FloatTensor(pos).unsqueeze_(1).to(args.device)
            pre = model(seq)

            adam_optimizer.zero_grad()
            loss = bce_criterion(pre, pos)

            # for param in model.item_emb0.parameters(): loss += args.l2_emb * torch.norm(param)
            # for param in model.item_emb1.parameters(): loss += args.l2_emb * torch.norm(param)
            # for param in model.item_emb2.parameters(): loss += args.l2_emb * torch.norm(param)
            # for param in model.item_emb3.parameters(): loss += args.l2_emb * torch.norm(param)
            # for param in model.item_emb4.parameters(): loss += args.l2_emb * torch.norm(param)
            # for param in model.item_emb5.parameters(): loss += args.l2_emb * torch.norm(param)
            # for param in model.item_emb6.parameters(): loss += args.l2_emb * torch.norm(param)
            # for param in model.item_emb7.parameters(): loss += args.l2_emb * torch.norm(param)
            # for param in model.item_emb8.parameters(): loss += args.l2_emb * torch.norm(param)
            # for param in model.item_emb9.parameters(): loss += args.l2_emb * torch.norm(param)

            loss.backward()
            adam_optimizer.step()
            # print("loss in epoch {} iteration {}: {}".format(epoch, step, loss.item())) # expected 0.4~0.6 after init few epochs
        if epoch % 1 == 0:
            model.eval()
            t1 = time.time() - t0
            T += t1
            print('Evaluating', end='')
            with torch.no_grad():
                test_auc, test_acc, test_recall = evaluate(model, user_test, args)
                train_auc, train_acc, train_recall = evaluate(model, user_train, args)
            print('epoch:%d, time: %f(s), valid (AUC: %.4f,ACC: %.4f,Recall: %.4f), test (AUC: %.4f,ACC: %.4f,Recall: %.4f)'
                    % (epoch, T, train_auc, train_acc, train_recall, test_auc, test_acc, test_recall))

            t_auc.append(test_auc)
            t_acc.append(test_acc)
            t_epoch.append(epoch)
    
            f.write(str(train_auc) + ' ' + str(test_auc) + '\n')
            f.flush()
            t0 = time.time()
            model.train()
    
        if epoch == args.num_epochs:
            folder = args.dataset + '_' + args.train_dir
            fname = 'SASRec.epoch={}.lr={}.layer={}.head={}.hidden={}.pth'
            fname = fname.format(args.num_epochs, args.lr, args.num_blocks, args.num_heads, args.hidden_units)
            torch.save(model.state_dict(), os.path.join(folder, fname))

    print('测试集最大AUC{}'.format(pd.DataFrame(t_auc).max()))
    print('测试集最大ACC{}'.format(pd.DataFrame(t_acc).max()))
    f.close()
    sampler.close()
    print("Done")
    auc_print(t_epoch, t_auc)
